from __future__ import print_function
import csv
import sys

'''
   Aims to provide an easy way to convert a DASH zmatrix to one that can be read by the qNRM code.
   Contact markspillman *at* gmail *dot* com for any queries.
   
   USAGE:
   Windows:	python.exe zmatrix.py input.zmatrix
   Linux:   python zmatrix.py input.zmatrix
'''

#Specify zmatrix for conversion as a command line argument.
filename = sys.argv[1]

# Read in the z-matrix, and create output file with same name, but in .txt format  --  in: file.zmatrix out: file.zmatrix.txt
in_zmatrix = csv.reader(open(filename, 'rb'), delimiter=' ')
out_zmatrix = csv.writer(open(filename+'.txt', 'wb'), delimiter=' ', quotechar=' ', quoting=csv.QUOTE_MINIMAL)


# This loop iterates over lines in the input zmatrix and changes certain values to match the required format for qNRM code.
nl = 0 #number of lines counter
matrix_str = [] # Stores the matrix
line_str = [] # Stores line components
for line in in_zmatrix:
	nl = nl + 1
	x = len(line) # Length of line
	if nl == 1:
		matrix_str.append('') # DASH outputs 'generated by xxxxxxx'. This is removed and replaced with a blank line.
	else:
		if nl == 2: # The second line has the unit cell dimensions.
			line[0] = raw_input('Enter cell dimension a: ') #DASH does not include these data. qNRM code requires it.
			line[1] = raw_input('Enter cell dimension b: ')
			line[2] = raw_input('Enter cell dimension c: ')
			line[3] = raw_input('Enter cell dimension al: ')
			line[4] = raw_input('Enter cell dimension be: ')
			line[5] = raw_input('Enter cell dimension ga: ')
			matrix_str.append(line)
			
		if nl == 3: # DASH outputs the number of atoms then a zero. qNRM requres a 1 followed by the number of atoms.
				temp = line[2] #Stores the number of atoms
				line = ['1', temp] #replaces the line with '1  N' where N is the number of atoms.
				matrix_str.append(line)
				
		if nl == 4:
				line[8] = '1.0000000' #Standard format for qNRM code. Creates dummy length and angle on first atom.
				line[15] = '90.00000'
				line[38] = '0.0000000'
				i = 39 
				while i < x: # Clear out last few columns
					line[i] = ''
					i = i+1
				matrix_str.append(line)
				
		if nl == 5:
				line[15] = '90.00000' # Dummy angle created on second atom.
				line[17] = '1' # Refinement flags for dummy.
				line[24] = '1' # Refinement flag.
				line[38] = '0.0000000'
				i = 39 
				while i < x: # Clear out last few columns
					line[i] = ''
					i = i+1
				matrix_str.append(line)
				
		if nl == 6:
				line[22] = '1' # Dummy refinement flag
				
		if nl > 5: # Clear out the last few columns, and replace with '0.0000000' - dummy.
			for item in line:
				if item != ' ':
					if item == '3.0':
						y = item.index('3.0') # Starts replacing at the Beq values, as the preceding column is the last one used in the qNRM code.
						y = y+1
						item = '0.0000000'
						while y < x:
							line[y] = ''
							y = y+1
						line_str.append(item)
					elif item == '6.0':
						y = item.index('6.0')
						y = y+1
						item = '0.0000000'
						while y < x:
							line[y] = ''
							y = y+1
						line_str.append(item)
					else:
						line_str.append(item)
			matrix_str.append(line_str) # Adds edited line to stored matrix
			line_str = [] # Clear stored line.
				

z = 0
while z != 1:
	i = raw_input('Enter spacegroup. For syntax help, enter ''-help'': ')
	if i == '-help':
		print('')
		print('Available spacegroups are:')
		print('P1, Pbar1, P2_1, P 1 21/c 1')
		print('')
	else:
		matrix_str.append([i])
		z = 1
	


out_zmatrix.writerows(matrix_str) # Outputs zmatrix to file.
print('')
print('qRNM zmatrix outputted as:',filename+'.txt')
print('DOUBLE CHECK Z-MATRIX BEFORE USE!')


